module(..., package.seeall)


set = class("set")
_G.set = set

function set.__init(self, ...)
    self.elements = {}
    self.size = 0
    if select("#", ...) > 0 then
        for _, elem in pairs{...} do
            self:add(elem)
        end
    end
end

function set.__tostring(self)
    local parts = {self.__name, "("}
    for elem in iter(self) do
        parts[#parts + 1] = tostring(elem)
        parts[#parts + 1] = ", "
    end
    -- remove the last comma if the set is non-empty
    if self.size > 0 then 
        parts[#parts] = nil 
    end
    parts[#parts + 1] = ")"
    return table.concat(parts)
end

function set.__len(self)
    return self.size
end

function set.__iter(self)
    return next, self.elements, nil
end
set.iter = set.__iter

function set.copy(self, target)
    target = target or set()
    for elem in iter(self) do
        target:add(elem)
    end
    return target
end

function set.add(self, elem)
    if not self.elements[elem] then
        self.elements[elem] = true
        self.size = self.size + 1
    end
end

function set.remove(self, elem)
    if not self.elements[elem] then
        error("element not found in set", 2)
    end
    self.elements[elem] = nil
    self.size = self.size - 1
end

function set.discard(self, elem)
    if self.elements[elem] then
        self:remove(elem)
    end
end

function set.pop(self)
    if self.size == 0 then
        error("cannot pop() from empty set", 2)
    end
    local elem = next(self.elements)
    self:remove(elem)
    return elem
end

function set.clear(self)
    self.elements = {}
    self.size = 0
end

function set.intersect(self, other)
    --[[ In-place set intersection operation. ]]
    for elem in iter(self) do
        if not other:contains(elem) then
            self:remove(elem)
        end
    end
    return self
end

function set.include(self, other)
    --[[ In-place set union operation. ]]
    for elem in iter(other) do
        self:add(elem)
    end
    return self
end

function set.exclude(self, other)
    --[[ In-place set difference operation. ]]
    for elem in iter(other) do
        self:discard(elem)
    end
    return self
end

function set.intersection(self, other)
    return self:copy():intersect(other)
end

function set.union(self, other)
    return self:copy():include(other)
end

function set.difference(self, other)
    return self:copy():exclude(other)
end

function set.symmetric_difference(self, other)
    return self:union(other):exclude(self:intersection(other))
end

function set.contains(self, elem)
    return self.elements[elem] ~= nil
end

function set.is_disjoint(self, other)
    for elem in iter(self) do
        if other:contains(elem) then
            return false
        end
    end
    return true
end
    
function set.is_subset(self, other)
    for elem in iter(self) do
        if not other:contains(elem) then
            return false
        end
    end
    return true
end

function set.is_proper_subset(self, other)
    return self.size < other.size and self:is_subset(other)
end

function set.is_superset(self, other)
    return other:is_subset(self)
end

function set.is_proper_superset(self, other)
    return other:is_proper_subset(self)
end

function set.equals(self, other)
    return self.size == other.size and self:is_subset(other)
end

set.__eq = set.equals
set.__le = set.is_subset
set.__lt = set.is_proper_subset
set.__add = set.union
set.__mul = set.intersection
set.__sub = set.difference
set.__pow = set.symmetric_difference


function test()
    local s = set(1,2,3,4,5)
    print(s)
    assert(s:contains(2))
    assert(not s:contains(10))
    
    local t = set(2, 3)
    print(t)
    assert(not(s < t))
    assert(not(s <= t))
    assert(not(s == t))
    assert(s >= t)
    assert(s > t)
    
    local u = s + t
    print(u)
    assert(u == s)

    local v = set(3,4,5,6,7,8,9,10) 
    print(v)
    assert(u + v == set(1,2,3,4,5,6,7,8,9,10))
    assert(u * v == set(3,4,5))
    assert(u - v == set(1,2))
    assert(v - u == set(6,7,8,9,10))
    assert(u ^ v == (u + v) - (u * v))
    print(u ^ v)
    return s, t, u, v
end
